export const createModel = () => {
  // 实例化模型
  /**
   * sequential
   */
  const model = tf.sequential();

  // 添加单行输入图层
  model.add(
    tf.layers.dense({
      inputShape: [1],
      units: 1, // 权重
      useBias: true,
    })
  );

  //  添加输出图层
  model.add(tf.layers.dense({ units: 1, useBias: true }));

  return model;
};

/**
 * 数据转化为 张量
 * @param {*} data
 * @returns
 */
export const convertToTensor = (data) => {
  return tf.tidy(() => {
    // 第一步：重新排列数据
    tf.util.shuffle(data);

    // 第二步： 数据 转化为 张量
    const inputs = data.map((d) => d.horsepower);
    const labels = data.map((d) => d.mpg);

    const inputTensor = tf.tensor2d(inputs, [inputs.length, 1]);
    const labelTensor = tf.tensor2d(labels, [labels.length, 1]);

    // 第三步： 对数据进行归一化
    const inputMax = inputTensor.max();
    const inputMin = inputTensor.min();
    const labelMax = labelTensor.max();
    const labelMin = labelTensor.min();

    const normalizedInputs = inputTensor
      .sub(inputMin)
      .div(inputMax.sub(inputMin));
    const normalizedLabels = labelTensor
      .sub(labelMin)
      .div(labelMax.sub(labelMin));

    return {
      inputs: normalizedInputs,
      labels: normalizedLabels,
      //
      inputMax,
      inputMin,
      labelMax,
      labelMin,
    };
  });
};

export const trainModel = async (model, inputs, labels) => {
  // 编译模型
  model.compile({
    optimizer: tf.train.adam(),
    loss: tf.losses.meanSquaredError,
    metrics: ["mse"]
  })

  const batchSize = 32
  const epochs = 50

  /**
   * 启动模型循环
   */
  return await model.fit(inputs, labels, {
    batchSize,
    epochs,
    shuffle: true,
    callbacks: tfvis.show.fitCallbacks(
      {name: "模型训练绩效"},
      ["loss", "mse"],
      { height: 200, callbacks: ["onEpochEnd"] }
    )
  })
};
